/*++

Copyright (c)  1999 - 2015 Intel Corporation. All rights reserved
This software and associated documentation (if any) is furnished
under a license and may only be used or copied in accordance
with the terms of the license. Except as permitted by such
license, no part of this software or documentation may be
reproduced, stored in a retrieval system, or transmitted in any
form or by any means without the express written consent of
Intel Corporation.


Module Name:

  S3NvramSave.c


--*/

#include "S3NvramSave.h"
#include <Protocol\PlatformType.h>
#include <Protocol/Decompress.h>
#include <Library/CompressLib.h>
//#include <Protocol/VariableLock.h> // AptioV server override: Commented protocol not installed in Core 5.011
#ifndef EFI_VARIABLE_RUNTIME_ALLOW_RECLAIM
#define EFI_VARIABLE_RUNTIME_ALLOW_RECLAIM 0x00001000
#endif


VOID
SaveS3StructToNvram (
  VOID
  )
/*++

Routine Description:

  This function will retrieve the S3 data from HOBs produced by MRC
  and will save it to NVRAM if the data is absent or different from
  the previously saved data.

Arguments:

  VOID

Returns:

  VOID

--*/
{
  EFI_STATUS                      Status;
  UINTN                           BufferSize;
  UINTN                           CurrentHobSize;
  UINTN                           RemainingHobSize;
  UINTN                           S3ChunkSize;
  CHAR16                          EfiMemoryConfigVariable[] = L"MemoryConfig0";
  EFI_HOB_GUID_TYPE               *GuidHob      = NULL;
  VOID                            *HobData      = NULL;
  VOID                            *VariableData = NULL;

  UINTN                           CompressedDataSize;
  UINT32                          ScratchSize;
  VOID                            *CompressedData = NULL;
  VOID                            *Scratch = NULL;
  EFI_DECOMPRESS_PROTOCOL         *Decompress = NULL;
  VOID                            *CompressedVariableData = NULL;
  UINTN                           CompressedBufferSize;
//  EDKII_VARIABLE_LOCK_PROTOCOL    *VariableLock = NULL; AptioV server override: Commented protocol not installed in Core 5.011

  RemainingHobSize = sizeof (struct sysNvram);

  DEBUG((EFI_D_INFO, "SaveS3StructToNvram: sysNvram size to be saved to NVRAM(bytes): %d\n", RemainingHobSize));

  //
  // Get first S3 data HOB
  //
  GuidHob = GetFirstGuidHob (&gEfiMemoryConfigDataHobGuid);
  // AptioV Server Override - Start: Return if gEfiMemoryConfigData Hob is not found
  if (GuidHob == NULL) {
    return ;
  }
  // AptioV Server Override - End: Return if gEfiMemoryConfigData Hob is not found

  Status = gBS->LocateProtocol (&gEfiDecompressProtocolGuid, NULL, (VOID **) &Decompress);
  if(EFI_ERROR (Status)) {
    DEBUG((EFI_D_INFO, "Decompress locate protocol error\n"));
    ASSERT_EFI_ERROR (Status);
    return;
  }

  //AptioV server override: Commented protocol not installed in Core 5.011
  //Status = gBS->LocateProtocol(&gEdkiiVariableLockProtocolGuid, NULL, (VOID **)&VariableLock);
  //ASSERT_EFI_ERROR(Status); 

  while (RemainingHobSize) {

    HobData = GET_GUID_HOB_DATA (GuidHob);
    CurrentHobSize = GET_GUID_HOB_DATA_SIZE (GuidHob);

    //
    // Trim the Hob data padding so there is no overflow on the buffer
    //
    if(CurrentHobSize > RemainingHobSize){
      CurrentHobSize = RemainingHobSize;
    }

    if (CurrentHobSize < MAX_HOB_ENTRY_SIZE){
      RemainingHobSize = 0;
    } else {
      RemainingHobSize -= CurrentHobSize;
    }

    //
    // Use the HOB data to save Memory Configuration Data
    //
    BufferSize = CurrentHobSize;
    Status = gBS->AllocatePool ( EfiBootServicesData,
                      BufferSize,
                      (VOID**)&VariableData
                      );

    ASSERT (VariableData != NULL);
    S3ChunkSize = MAX_HOB_ENTRY_SIZE / 8;
    while (CurrentHobSize) {
      if (S3ChunkSize > CurrentHobSize) {
        S3ChunkSize = CurrentHobSize;
      }
      BufferSize = S3ChunkSize;
      CompressedDataSize = 0;
      ScratchSize = 0;
      Status = gRT->GetVariable (
                      EfiMemoryConfigVariable,
                      &gEfiMemoryConfigDataGuid,
                      NULL,
                          &CompressedDataSize,
                          NULL
                          );
      DEBUG((DEBUG_INFO, "MemoryConfig variable result=%r, size=0x%x\n", Status, CompressedDataSize));

      if(Status == EFI_BUFFER_TOO_SMALL) {
        Status = gBS->AllocatePool ( EfiBootServicesData,
                        CompressedDataSize,
                        (VOID**)&CompressedData
                        );
        ASSERT (Status == EFI_SUCCESS); 
      }

      if(!EFI_ERROR (Status))
      {
        Status = gRT->GetVariable (
                         EfiMemoryConfigVariable,
                         &gEfiMemoryConfigDataGuid,
                         NULL,
                         &CompressedDataSize,
                         CompressedData
                         );

        if (!EFI_ERROR (Status)) {
         Status = Decompress->GetInfo (
                         Decompress,
                         CompressedData,
                         (UINT32)CompressedDataSize,
                         &(UINT32)BufferSize,
                         &ScratchSize
                         );
        }

        if (!EFI_ERROR (Status)) {
          Status = gBS->AllocatePool ( EfiBootServicesData,
                         ScratchSize,
                         (VOID**)&Scratch
                         );
        }

        if (!EFI_ERROR (Status)) {
          Status = Decompress->Decompress (
                          Decompress,
                          CompressedData,
                          (UINT32)CompressedDataSize,
                           VariableData,
                           (UINT32)BufferSize,
                           Scratch,
                           ScratchSize
                           );
        }

        if (EFI_ERROR (Status)) {
          DEBUG((EFI_D_ERROR, "Getting variables error: 0x%x\n", Status));
          ASSERT (Status == EFI_SUCCESS); 
        }

        if(Scratch != NULL) {
          gBS->FreePool (Scratch);
          Scratch = NULL;
        }
      }

      if(CompressedData != NULL) {
        gBS->FreePool (CompressedData);
        CompressedData = NULL;
      }

      if ( (EFI_ERROR(Status)) || (CompareMem (HobData, VariableData, S3ChunkSize) != 0) ) {
        Status = gBS->AllocatePool ( EfiBootServicesData,
                          BufferSize,
                          (VOID**)&CompressedVariableData
                          );
        ASSERT (CompressedVariableData != NULL); 
        if (Status == EFI_SUCCESS) {
          CompressedBufferSize = BufferSize;
          Status = Compress(HobData, S3ChunkSize, CompressedVariableData, &CompressedBufferSize);

          if(Status == EFI_SUCCESS) {
            Status = gRT->SetVariable (
                           EfiMemoryConfigVariable,
                           &gEfiMemoryConfigDataGuid,
                           EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS, //AptioV Server Override: Remove attribute not supported by AMI NVRAM driver.
                           CompressedBufferSize,
                           CompressedVariableData
                           );
          }
          if(CompressedVariableData != NULL) {
            gBS->FreePool(CompressedVariableData);
            CompressedVariableData = NULL;
          }
          DEBUG((EFI_D_INFO, "SaveS3StructToNvram: Variable#: %c, Set variable status: %r \n",  EfiMemoryConfigVariable[12], Status));
        }

        if (EFI_ERROR (Status)) {
          DEBUG((EFI_D_ERROR, "Set variable error. Status: 0x%x\n", Status));
          ASSERT_EFI_ERROR (Status);
        }
      }
/*  AptioV server override: Commented protocol not installed in Core 5.011
      //
      // Lock the Memory Config Variable
      //
      if (!EFI_ERROR(Status)) {
        Status = VariableLock->RequestToLock(VariableLock, EfiMemoryConfigVariable, &gEfiMemoryConfigDataGuid);
        ASSERT_EFI_ERROR(Status);
      }  
*/  
      HobData = (UINT8 *) (HobData) + S3ChunkSize;
      CurrentHobSize -= S3ChunkSize;
      EfiMemoryConfigVariable[12]++;  // Increment number in the string
    }
    //
    // Get next S3 Config data hob, if none left, results NULL
    //
    GuidHob = GET_NEXT_HOB (GuidHob); // Increment to next HOB
    GuidHob = GetNextGuidHob (&gEfiMemoryConfigDataHobGuid, GuidHob);  // Now search for next MemConfig HOB
    if(VariableData != NULL) {
        gBS->FreePool(VariableData);
        VariableData = NULL;
    }
  }
  return;
}


BOOLEAN
S3NvramCompareStructs (
  IN VOID                               *CurrentBootStruct,
  IN VOID                               *SavedBootStruct
  )
/*++

Routine Description:

  This function is used to compare the saved S3 structure from NVRAM
  with the S3 structure populated on the current boot.  This will
  determine if the structure needs to be saved to NVRAM again or not.

Arguments:

  IN VOID  *CurrentBootStruct - Pointer to current boot NVRAM structure
  IN VOID  *SavedBootStruct - Pointer to buffer of the saved NVRAM structure
                              from a previous boot

Returns:

  BOOLEAN  -  TRUE  - The structures match
              FALSE - There are differences

--*/
{
  EFI_STATUS                    Status;
  struct sysNvram               *SavedS3    = SavedBootStruct;
  struct sysNvram               *CurrentS3  = CurrentBootStruct;
  UINTN                         CompareValue;
  EFI_PLATFORM_TYPE_PROTOCOL    *mPlatformType;
  BOOLEAN                       NotSthiPlatform = TRUE;
  BOOLEAN                       StructsMatch = TRUE;

  Status = gBS->LocateProtocol (
                &gEfiPlatformTypeProtocolGuid,
                NULL,
                &mPlatformType
                );
  ASSERT_EFI_ERROR (Status);

  //
  // TODO:  This will be updated to check more vital parts of the
  // NVRAM structure instead of the entire structure.  For now, a compare of
  // the entire NVRAM structure is used
  //
  CompareValue = CompareMem ((VOID*)CurrentS3, (VOID*)SavedS3, sizeof (struct sysNvram));

#ifndef GRANGEVILLE_FLAG
  //
  // For AztecCitySTHI, skip saving this structure a 2nd time to avoid garbage collection.  Config never changes
  //
  if (mPlatformType->BoardId == TypeAztecCitySTHI)
    NotSthiPlatform = FALSE;
#endif

  if((CompareValue != 0) && NotSthiPlatform)
    StructsMatch = FALSE;

  return StructsMatch;
}


EFI_STATUS
EFIAPI
S3NvramSaveEntry (
  IN EFI_HANDLE                         ImageHandle,
  IN EFI_SYSTEM_TABLE                   *SystemTable
  )
/*++

Routine Description:

  This is the main entry point of the S3 NVRAM Save module.

Arguments:

  ImageHandle  -  Handle for the image of this driver.
  SystemTable  -  Pointer to the EFI System Table.

Returns:

  EFI_SUCCESS  -  Module launched successfully.

--*/
{
  EFI_STATUS Status = EFI_SUCCESS;

  //
  // Save the s3 strututre from MRC into NVRAM if needed
  //
  SaveS3StructToNvram();

  return Status;

}
